#!/bin/bash

# Copyright (c) Microsoft Corporation
# All rights reserved.
#
# MIT License
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
# documentation files (the "Software"), to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
# to permit persons to whom the Software is furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED *AS IS*, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
# BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.


# Bootstrap script for docker container.

exec 17>/pai/log/DockerContainerDebug.log
BASH_XTRACEFD=17

function exit_handler()
{
  printf "%s %s\n" \
    "[DEBUG]" "Docker container exit handler: EXIT signal received in docker container, exiting ..."
}

set -x
PS4="+[\t] "
trap exit_handler EXIT


touch "/alive/docker_$PAI_CONTAINER_ID"



export PAI_WORK_DIR="$(pwd)"
PAI_WEB_HDFS_PREFIX={{{ webHdfsUri }}}/webhdfs/v1/Container
HDFS_LAUNCHER_PREFIX=$PAI_DEFAULT_FS_URI/Container
export CLASSPATH="$(hadoop classpath --glob)"

task_role_no={{{ idx }}}

printf "%s %s\n%s\n\n" "[INFO]" "ENV" "$(printenv | sort)"

cp -r /pai/code/* ./

function webhdfs_create_file()
{
  webHdfsRequestPath=${1}"?user.name="{{{ jobData.userName }}}"&op=CREATE"
  redirectResponse=$(curl -i -X PUT ${webHdfsRequestPath} -o /dev/null -w %{redirect_url}' '%{http_code})
  redirectCode=$(cut -d ' ' -f 2 <<< ${redirectResponse})
  if [[ ${redirectCode} = "307" ]]; then
    redirectUri=$(cut -d ' ' -f 1 <<< ${redirectResponse})
    createResponse=$(curl -i -S -X PUT ${redirectUri})
  else
    printf "%s %s\n %s %s\n %s %s\n" \
      "[WARNING]" "Webhdfs creates folder failed" \
      "Folder Path:" ${webHdfsRequestPath} \
      "Response code:" ${redirectCode}
  fi
}

function webhdfs_download_file()
{
  webHdfsRequestPath=${1}"?user.name="{{{ jobData.userName }}}"&op=OPEN"
  localPath=${2}
  downloadResponse=$(curl -S -L ${webHdfsRequestPath} -o ${localPath} -w %{http_code})
  if [[ ${downloadResponse} = "200" ]]; then
    printf "%s %s\n" \
      "[INFO]" "Webhdfs downloads file succeed"
  else
    printf "%s %s\n" \
      "[WARNING]" "Webhdfs downloads file failed"
  fi
}

function prepare_ssh()
{
  mkdir /root/.ssh
  sed -i 's/PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config
  sed 's@session\s*required\s*pam_loginuid.so@session optional pam_loginuid.so@g' -i /etc/pam.d/sshd
}

function start_ssh_service()
{
  printf "%s %s\n" \
    "[INFO]" "start ssh service"
  cat /root/.ssh/{{{ jobData.jobName }}}.pub >> /root/.ssh/authorized_keys
  sed -i 's/Port.*/Port '$PAI_CONTAINER_SSH_PORT'/' /etc/ssh/sshd_config
  echo "sshd:ALL" >> /etc/hosts.allow
  service ssh restart
}

function get_ssh_key_files()
{
  info_source="webhdfs"
  localKeyPath=/root/.ssh/{{{ jobData.jobName }}}.pub
  localPrivateKeyPath=/root/.ssh/{{{ jobData.jobName }}}

  if [[ -f $localKeyPath ]]; then
    rm -f $localKeyPath
  fi

  if [[ "$info_source" = "webhdfs" ]]; then
    webHdfsKeyPath=${PAI_WEB_HDFS_PREFIX}/{{{ jobData.userName }}}/{{{ jobData.jobName }}}/ssh/keyFiles/{{{ jobData.jobName }}}.pub
    webhdfs_download_file $webHdfsKeyPath $localKeyPath
    webHdfsPrivateKeyPath=${PAI_WEB_HDFS_PREFIX}/{{{ jobData.userName }}}/{{{ jobData.jobName }}}/ssh/keyFiles/{{{ jobData.jobName }}}
    webhdfs_download_file $webHdfsPrivateKeyPath $localPrivateKeyPath
    chmod 400 ${localPrivateKeyPath}
  else
    printf "%s %s\n" \
      "[WARNING]" "Get another key store way"
  fi
}

function generate_ssh_connect_info()
{
  info_source="webhdfs"
  destFileName=${1}

  if [[ "$info_source" = "webhdfs" ]]; then
    webHdfsRequestPath=$destFileName
    webhdfs_create_file $webHdfsRequestPath
  else
    printf "%s %s\n" \
      "[WARNING]" "Get another key store way"
  fi
}

# Check whether hdfs bianry and ssh exists, if not ignore ssh preparation and start part
# Start sshd in docker container
if service --status-all 2>&1 | grep -q ssh; then
  prepare_ssh
  get_ssh_key_files
  sshConnectInfoFolder=${PAI_WEB_HDFS_PREFIX}/${PAI_USER_NAME}/${PAI_JOB_NAME}/ssh/$APP_ID
  # Generate ssh connect info file in "PAI_CONTAINER_ID-PAI_CURRENT_CONTAINER_IP-PAI_CONTAINER_SSH_PORT" format on hdfs
  destFilePath=${sshConnectInfoFolder}/$PAI_CONTAINER_ID-$PAI_CONTAINER_HOST_IP-$PAI_CONTAINER_SSH_PORT
  generate_ssh_connect_info ${destFilePath}

  # Enable SSH connect inside
  # Move ssh config file to /root/.ssh/
  mv /pai/ssh_config/config /root/.ssh/
  # Start ssh service
  start_ssh_service
fi

# Write env to system-wide environment
env | grep -E "^PAI|PATH|PREFIX|JAVA|HADOOP|NVIDIA|CUDA" > /etc/environment

function run_user_command()
{
  printf "%s %s\n\n" "[INFO]" "USER COMMAND START"
  {{{ taskData.command }}} || exit $?
  printf "\n%s %s\n\n" "[INFO]" "USER COMMAND END" # job_exporter relay on this exact output, if you are going to change this, remember to change job_exporter also
  exit 0
}

run_user_command &
user_command_pid=$!

while [ $(( $(date +%s) - $(stat -c %Y /alive/yarn_$PAI_CONTAINER_ID) )) -lt 30 ] && \
        kill -0 $user_command_pid 2>/dev/null; do
  sleep 20
done

if kill -0 $user_command_pid 2>/dev/null; then
  echo "job has been killed, docker container exiting"
  exit 0
else
  wait $user_command_pid
  user_command_exitcode=$?
  echo "job has finished with exit code $user_command_exitcode"
  exit $user_command_exitcode
fi

